#! /usr/bin/python3
from enum import Enum
import copy
import time
import colorsys
import os,sys,argparse,random
from timeit import default_timer as timer
import cv2
import numpy as np
from keras import backend as K
from keras.models import load_model
from keras.layers import Input
from PIL import Image, ImageFont, ImageDraw

from nets.yolo4_tiny import yolo_body, yolo_eval
from utils.utils import letterbox_image

from yolo import YOLO

from tqdm import tqdm
from scipy import misc

from matplotlib import pyplot as plt

from KalmanFilterTracker import Tracker  # 加载卡尔曼滤波函数
import math


def calc_center(out_boxes,out_classes,out_scores,score_limit = 0.5):
    outboxes_filter = []
    for x,y,z in zip(out_boxes,out_classes,out_scores):
        if z > score_limit:
            if y == 0: # person标签为0
                outboxes_filter.append(x)
    
    centers= []
    number = len(outboxes_filter)
    for box in outboxes_filter:
        top, left, bottom, right = box
        center=np.array([[(left+right)//2],[(top+bottom)//2]])
        centers.append(center)
    return centers,number


def get_colors_for_classes(num_classes):
    """Return list of random colors for number of classes given."""
    # Use previously generated colors if num_classes is the same.
    if (hasattr(get_colors_for_classes, "colors") and
            len(get_colors_for_classes.colors) == num_classes):
        return get_colors_for_classes.colors

    hsv_tuples = [(x / num_classes, 1., 1.) for x in range(num_classes)]
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(
        map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
            colors))
    #colors = [(255,99,71) if c==(255,0,0) else c for c in colors ]  # 单独修正颜色，可去除
    random.seed(10101)  # Fixed seed for consistent colors across runs.
    random.shuffle(colors)  # Shuffle colors to decorrelate adjacent classes.
    random.seed(None)  # Reset seed to default.
    get_colors_for_classes.colors = colors  # Save colors for future calls.
    return colors

def cross_mul(vec1, vec2):
    x1, y1 = vec1
    x2, y2 = vec2
    return x1*y2 - x2*y1

def judge_three_points(p1, p2, p3):
    x1, y1 = p1
    x2, y2 = p2
    x3, y3 = p3
    cross_mul_val = cross_mul((x2-x1, y2-y1), (x3-x1, y3-y1))
    if(cross_mul_val >0):
        return 1
    elif(cross_mul_val < 0):
        return -1
    return 0

def judge_line(p1, p2, delta_base):
    x1, y1 = p1
    x2, y2 = p2

    delta_x = abs(x1 - x2)
    delta_y = abs(y2 - y1)

    if (delta_y > delta_x) and (delta_y > delta_base):
        if(y2 > y1):
            return 3
        else:
            return 4
    
    if (delta_x > delta_y) and (delta_x > delta_base):
        if(x2 > x1):
            return 2
        else:
            return 1

    return 0




def judge_direction(ppoints):
    """
    @function: judge the hand move direction according to centers
    @ppoints: centers
    @outputs: direction: 
                0: no direction
                1: right; 2: left; 3: up; 4: down
                5: clockwise; 6: counter_clockwise
    """
    judge_base = 15
    tolerance = 100 

    N = len(ppoints)
    x1, y1 = ppoints[0]
    x_last, y_last = ppoints[N-1]
    if (abs(x_last - x1) < judge_base) and (abs(y_last - y1) < judge_base):
        return 0

    ## judge clockwise and counter_clockwise
    tmp_res = 0
    for i in range(0, N-2):
        tmp_res += judge_three_points(ppoints[i], ppoints[i+1], ppoints[i+2])

    if (tmp_res > 0) and (tmp_res > (N-2)*0.8):
        return 5 ## clockwise
    elif (tmp_res < 0) and (abs(tmp_res) > (N-2)*0.8):
        return 6 ## counter_clockwise

    line_res = [0] * 4;
    for i in range(0, N-1, 2):
        xi, yi = ppoints[i]
        xi_next, yi_next = ppoints[i+1]

        if (i == 0):
            if(xi_next >= xi) and (abs(yi_next - yi) <= abs(xi_next - xi)):
                line_res[0] = 1 ## right

            if (xi_next <= xi) and (abs(yi_next - yi) <= abs(xi_next - xi)):
                line_res[1] = 1 ## left

            if (yi_next >= yi) and (abs(xi_next - xi) <= abs(yi_next - yi)):
                line_res[2] = 1 ## up

            if (yi_next <= yi) and (abs(xi_next - xi) <= abs(yi_next - yi)):
                line_res[3] = 1 ## down
        else:
            if (line_res[0] == 1) and ((xi_next < xi) or (abs(yi_next - yi) > (xi_next - xi))):
                line_res[0] = 0

            if (line_res[1] == 1) and ((xi_next > xi) or (abs(yi_next - yi) > (xi_next - xi))):
                line_res[1] = 0

            if (line_res[2] == 1) and ((yi_next < yi) or (abs(xi_next - xi) > (yi_next - yi))):
                line_res[2] = 0

            if  (line_res[3] == 1) and ((yi_next > yi) or (abs(xi_next - xi) > (yi_next - yi))):
                line_res[3] = 0

    for i in range(4):
        if(line_res[i] == 1):
            return i+1

    return 0

def trackerDetection(tracker,image,centers,number,max_point_distance = 30,max_colors = 20,track_id_size = 1):
    '''
        - max_point_distance为两个点之间的欧式距离不能超过30
            - 有多条轨迹,tracker.tracks;
            - 每条轨迹有多个点,tracker.tracks[i].trace
        - max_colors,最大颜色数量
        - track_id_size,每个中心点，显示的标记字体大小
    '''
    direction_y_base = 100
    direction_font_scale = 1.0
    direction_font_color = (51, 255, 51)
    direction_thickness = 2
    track_colors = get_colors_for_classes(max_colors)
    
    result = np.asarray(image)
    font = cv2.FONT_HERSHEY_SIMPLEX
    #cv2.putText(result, str(number)+'hand', (20,  40), font, 0.8, (190, 164, 137), 1)  # 左上角的hand计数

    if (len(centers) > 0):
        # Track object using Kalman Filter
        tracker.Update(centers)
        # For identified object tracks draw tracking line
        # Use various colors to indicate different track_id
        for i in range(len(tracker.tracks)):
            # 多个轨迹
            hand_gesture = 0
            if (len(tracker.tracks[i].trace) > 1):
                x0,y0 = tracker.tracks[i].trace[-1][0][0],tracker.tracks[i].trace[-1][1][0]
                cv2.putText(result,str(tracker.tracks[i].track_id),(int(x0),int(y0)),font,track_id_size,(255, 255, 255),2)  

                # (image,text,(x,y),font,size,color,粗细)
                trace_size = len(tracker.tracks[i].trace)
                ppoints = []
                for j in range(0, trace_size):
                    #每条轨迹的每个点
                    # Draw trace line
                    x1 = tracker.tracks[i].trace[j][0][0]
                    y1 = tracker.tracks[i].trace[j][1][0]
                                      
                    clr = tracker.tracks[i].track_id % 9
                    cv2.circle(result, (int(x1), int(y1)), 3, track_colors[clr], 3)

                    ppoints.append((x1, y1))

                hand_gesture = judge_direction(ppoints)

                ## put hand movement info on picture
                if(hand_gesture == 1):
                    cv2.putText(result, 'Hand'+str(i+1)+': Right', (10, direction_y_base+i*50), font, direction_font_scale, direction_font_color, direction_thickness) 
                elif(hand_gesture == 2):
                    cv2.putText(result, 'Hand'+str(i+1)+': Left', (10, direction_y_base+i*50), font, direction_font_scale, direction_font_color, direction_thickness) 
                elif(hand_gesture == 3):
                    cv2.putText(result, 'Hand'+str(i+1)+': Up', (10, direction_y_base+i*50), font, direction_font_scale, direction_font_color, direction_thickness) 
                elif(hand_gesture == 4):
                    cv2.putText(result, 'Hand'+str(i+1)+': Down', (10, direction_y_base+i*50), font, direction_font_scale, direction_font_color, direction_thickness) 
                elif(hand_gesture == 5):
                    cv2.putText(result, 'Hand'+str(i+1)+': Clockwise', (10, direction_y_base+i*50), font, direction_font_scale, direction_font_color, direction_thickness) 
                elif(hand_gesture == 6):
                    cv2.putText(result, 'Hand'+str(i+1)+': CounterClockwise', (10, direction_y_base+i*50), font, direction_font_scale, direction_font_color, direction_thickness) 

    return tracker,result



# yolo 配置
yolo_args = {
    "model_path": 'logs/ep084-loss1.536-val_loss0.611.h5',
    "anchors_path": 'model_data/yolo_anchors.txt',
    "classes_path": 'model_data/hand.txt',
    "score" : 0.4,
    "iou" : 0.3,
    "model_image_size" : (416, 416),
}


yolo = YOLO(**yolo_args)


tracker = Tracker(100, 8, 8, 1)   
#第一个数字是距离阈值，超过即删除，重画新的轨迹，默认100；
#第二个数字是当未被检测到时，允许跳过的最大的帧数,默认8；
#第三个数字是traceline的长度，默认15；
#第四个数字是hand的ID=100,101,102...

cap = cv2.VideoCapture(0)

cv2.namedWindow("hand_detection", cv2.WINDOW_NORMAL)
cv2.resizeWindow("hand_detection", 1024, 600)

while(True):
    ret, frame = cap.read()
    if frame is None:
        break
    frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB) # BGR --> RGB
    image = Image.fromarray(np.uint8(frame)) 

    r_image,out_boxes, out_scores, out_classes = yolo.detect_image(image)    
    centers,number = calc_center(out_boxes,out_classes,out_scores,score_limit = 0.6)
    tracker,result = trackerDetection(tracker, r_image, centers, number, max_point_distance = 80)

    result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR) # RGB --> BGR
    cv2.imshow("hand_detection",result)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

print('Down!')
